"""
算法核心步骤
- 1. 初始化词汇表：将语料分解为字符（或最小单元）
- 2. 统计字符对频率：遍历语料，统计所有相邻字符对的出现次数
- 3. 合并高频字符对：选择频率最高的字符对合并为新单元
- 4. 更新词汇表与语料表示：用新单元替换原字符对，重复步骤 2-3
- 5. 终止条件：达到预设词汇表大小或字符对频率为 0

相对于V1版算法，想要达到的目标：
- 1. 支持任意字符，而不仅仅是字节，中文按照字节拆分可能会导致乱码；
- 2. 支持任意字符对，而不仅仅是相邻字符对
"""
from collections import Counter, defaultdict


class BPEAlgorithm:
    
    def __init__(self, num_merges: int = 100):
        self.num_merges = num_merges
        # 字符
        self.chars: set[str] = set()
        # 存储合并规则的字典: (索引对) -> 新索引
        self.merges: dict[str, int] = {}
        # 词汇表: 索引 -> 对应的字节序列
        self.vocab: dict[int, str] = {}
    
    def train(self, texts: list[str]) -> dict[int, str]:
        # 1. 初始化词汇表（字符级）
        for text in texts:
            tmp = set(text.split())
            self.chars.update(tmp)
        print(self.chars)
        self.chars.remove(",")
        self.chars.remove(" ")
        
        # 2. 统计字符对频率
        text_words = [text.split() for text in texts]
        pairs: dict[tuple[str, str], int] = defaultdict(int)
        for words in text_words:
            for word1, word2 in zip(words[:-1], words[1:]):
                if word1 in self.chars and word2 in self.chars:
                    pairs[(word1, word2)] += 1
        # 过滤掉频率过低的字符对
        pairs = {k: v for k, v in pairs.items() if v >= 2}
        print("字符对频率:", pairs)

        # 3. 合并高频字符对
        for num_merge in range(self.num_merges):
            print("="*50, num_merge)
            cout_pairs = Counter(pairs)
            pair, freq = cout_pairs.most_common(1)[0]
            word1, word2 = pair
            self.merges[word1 + word2] = len(self.merges)
        return self.vocab
        

if __name__ == "__main__":
    num_merges = 100
    print("\n" + "="*100)
    text = "low, lower, newest, widest"
    vocab = BPEAlgorithm(num_merges).train([text])
    print({k: v.decode("utf8") for k, v in vocab.items() if k >= 256})

    print("\n" + "="*100)
    text = "hello world, 你好"
    vocab = BPEAlgorithm(num_merges).train([text])
    print({k: v.decode("utf8") for k, v in vocab.items() if k >= 256})

    print("\n" + "="*100)
    texts = [
        "low, lower, newest, widest",
        "lowest, higher, widest",
        "newest, higher, widest",
    ]
    vocab = BPEAlgorithm(num_merges).train(texts)
    print({k: v.decode("utf8") for k, v in vocab.items() if k >= 256})
